import json
import random
import pickle
from tqdm import tqdm

# 定义植物信息
PLANTS = {
    "Peashooter": {"cost": 100, "damage": 20},
    "Sunflower": {"cost": 50, "sun_production": 25}
}

# 定义僵尸信息
ZOMBIES = {
    "NormalZombie": {"health": 150, "speed": 1, "damage": 10}
}

# 定义可能的操作，允许在第0行和第1行种植
ACTIONS = [
    {"action": "plant", "plant_type": "Peashooter", "row": 0, "col": 0},
    {"action": "plant", "plant_type": "Peashooter", "row": 1, "col": 0},
    {"action": "plant", "plant_type": "Sunflower", "row": 0, "col": 0},
    {"action": "plant", "plant_type": "Sunflower", "row": 1, "col": 0},
    {"action": "wait"}
]

# 定义超参数
ALPHA = 1.1  # 学习率
GAMMA = 0.9  # 折扣因子
EPSILON = 0.1  # 探索率

# 初始化 Q-table
Q_TABLE = {}

def get_state_key(game_state):
    """
    将游戏状态转换为唯一的键，用于 Q-table。
    包含阳光数量、每行第0列的植物类型和是否有僵尸。
    """
    sun = game_state["sun"]
    rows = [0, 1]
    plant_types = {}
    has_zombie = {}
    for row in rows:
        plant_at_col0 = [plant for plant in game_state["plants"] if plant["row"] == row and plant["col"] == 0]
        plant_types[row] = plant_at_col0[0]["type"] if plant_at_col0 else "none"
        has_zombie[row] = any(zombie["row"] == row for zombie in game_state["zombies"])
    state_key = (sun, (plant_types[0], has_zombie[0]), (plant_types[1], has_zombie[1]))
    return state_key

def choose_action(state_key, epsilon=EPSILON):
    """
    根据 Q-table 使用 epsilon-greedy 策略选择动作。
    """
    if random.uniform(0, 1) < epsilon:
        return random.choice(ACTIONS)
    else:
        if state_key not in Q_TABLE:
            return random.choice(ACTIONS)
        q_values = [Q_TABLE.get((state_key, tuple(action.items())), 0) for action in ACTIONS]
        max_index = q_values.index(max(q_values))
        return ACTIONS[max_index]

def update_q_table(state_key, action, reward, next_state_key):
    """
    更新 Q-table。
    """
    current_q = Q_TABLE.get((state_key, tuple(action.items())), 0)
    next_q_values = [Q_TABLE.get((next_state_key, tuple(next_action.items())), 0) for next_action in ACTIONS] if next_state_key in Q_TABLE else [0] * len(ACTIONS)
    max_next_q = max(next_q_values)
    new_q = current_q + ALPHA * (reward + GAMMA * max_next_q - current_q)
    Q_TABLE[(state_key, tuple(action.items()))] = new_q

def simulate_step(game_state, action):
    """
    模拟游戏一步，更新状态并计算奖励。
    """
    new_game_state = {
        "sun": game_state["sun"],
        "plants": [plant.copy() for plant in game_state["plants"]],
        "zombies": [zombie.copy() for zombie in game_state["zombies"]]
    }
    reward = 0

    # 处理种植动作
    if action["action"] == "plant":
        plant_type = action["plant_type"]
        row = action["row"]
        col = action["col"]
        if any(plant["row"] == row and plant["col"] == col for plant in new_game_state["plants"]):
            reward -= 5  # 位置已有植物，给予惩罚
        else:
            cost = PLANTS[plant_type]["cost"]
            if new_game_state["sun"] >= cost:
                new_game_state["sun"] -= cost
                new_game_state["plants"].append({
                    "type": plant_type,
                    "row": row,
                    "col": col,
                    "health": 100
                })
                reward += 10  # 成功种植奖励
            else:
                reward -= 5  # 阳光不足惩罚
    elif action["action"] == "wait":
        pass

    # 僵尸移动
    for zombie in new_game_state["zombies"]:
        zombie["col"] -= ZOMBIES[zombie["type"]]["speed"]

    # 植物攻击僵尸
    for plant in new_game_state["plants"]:
        if plant["type"] == "Peashooter":
            for zombie in new_game_state["zombies"]:
                if zombie["row"] == plant["row"]:
                    zombie["health"] -= PLANTS["Peashooter"]["damage"]

    # 移除死亡的僵尸并计算奖励
    num_zombies_before = len(new_game_state["zombies"])
    new_game_state["zombies"] = [zombie for zombie in new_game_state["zombies"] if zombie["health"] > 0]
    zombies_killed = num_zombies_before - len(new_game_state["zombies"])
    reward += zombies_killed * 50  # 杀死僵尸的奖励

    # 僵尸攻击植物
    for zombie in new_game_state["zombies"]:
        for plant in new_game_state["plants"]:
            if zombie["row"] == plant["row"] and zombie["col"] == plant["col"]:
                plant["health"] -= ZOMBIES[zombie["type"]]["damage"]

    # 移除死亡的植物
    new_game_state["plants"] = [plant for plant in new_game_state["plants"] if plant["health"] > 0]

    # 向日葵每步产生阳光
    for plant in new_game_state["plants"]:
        if plant["type"] == "Sunflower":
            new_game_state["sun"] += PLANTS["Sunflower"]["sun_production"]

    return new_game_state, reward

# 初始游戏状态（增加僵尸数量以丰富训练数据）
game_state = {
    "sun": 200,
    "plants": [
        {"type": "Sunflower", "row": 0, "col": 0, "health": 100}
    ],
    "zombies": [
        {"type": "NormalZombie", "row": 0, "col": 9, "health": 150},  # 增加一行僵尸
        {"type": "NormalZombie", "row": 1, "col": 9, "health": 150}
    ]
}
game_state_json = json.dumps(game_state)

# 训练过程
EPISODES = 500  # 增加训练轮数
for episode in tqdm(range(EPISODES), desc="训练进度"):
    current_game_state = json.loads(game_state_json)
    total_reward = 0
    plant_count = 0
    step = 0
    while True:
        state_key = get_state_key(current_game_state)
        action = choose_action(state_key)
        next_game_state, reward = simulate_step(current_game_state, action)
        if action["action"] == "plant" and reward >= 10:  # 成功种植
            plant_count += 1
        next_state_key = get_state_key(next_game_state)
        update_q_table(state_key, action, reward, next_state_key)
        current_game_state = next_game_state
        total_reward += reward
        step += 1

        # 终止条件：无僵尸或达到最大步数
        if not current_game_state["zombies"] or step > 100:
            break

    print(f"第 {episode + 1} 轮: 总奖励 = {total_reward}, 种植植物次数 = {plant_count}")

# 保存 Q-table 到文件
with open('q_table.pkl', 'wb') as f:
    pickle.dump(Q_TABLE, f)
print("Q-table 已保存到 q_table.pkl")

import json
import pickle

# 加载 Q-table
with open('q_table.pkl', 'rb') as f:
    Q_TABLE = pickle.load(f)
print("Q-table 已加载")

# 测试初始状态
test_game_state = json.loads(game_state_json)
test_state_key = get_state_key(test_game_state)
next_action = choose_action(test_state_key, epsilon=0)  # epsilon=0 表示不探索
print("训练后的下一个动作:", next_action)